package main

import (
	"encoding/xml"
	"fmt"
	"os/exec"
)

type usedInfo struct {
	Total string `xml:"total"`
	Used  string `xml:"used"`
	Free  string `xml:"free"`
}

type fbMemoryUsage struct {
	usedInfo
}

type bar1MemoryUsage struct {
	usedInfo
}

// 温度相关字段
type temperature struct {
	GPUTemp                string `xml:"gpu_temp"`
	GPUTempMaxThreshold    string `xml:"gpu_temp_max_threshold"`
	GPUTempSlowThreshold   string `xml:"gpu_temp_slow_threshold"`
	GPUTempMaxGPUThreshold string `xml:"gpu_temp_max_gpu_threshold"`
	GPUTargetTemperature   string `xml:"gpu_target_temperature"`
	MemoryTemp             string `xml:"memory_temp"`
	GPUTempMaxMemThreshold string `xml:"gpu_temp_max_mem_threshold"`
}

type powerReadings struct {
	PowerState         string `xml:"power_state"`
	PowerManagement    string `xml:"power_management"`
	PowerDraw          string `xml:"power_draw"`
	PowerLimit         string `xml:"power_limit"`
	DefaultPowerLimit  string `xml:"default_power_limit"`
	EnforcedPowerLimit string `xml:"enforced_power_limit"`
	MinPowerLimit      string `xml:"min_power_limit"`
	MaxPowerLimit      string `xml:"max_power_limit"`
}

type clocks struct {
	GraphicsClock string `xml:"graphics_clock"`
	SMClock       string `xml:"sm_clock"`
	MEMClock      string `xml:"mem_clock"`
	VideoClock    string `xml:"video_clock"`
}

type gpuInfo struct {
	ID                  string          `xml:"id,attr"`
	ProductName         string          `xml:"product_name"`
	ProductBrand        string          `xml:"product_brand"`
	ProductArchitecture string          `xml:"product_architecture"`
	DisplayMode         string          `xml:"display_mode"`
	DisplayActive       string          `xml:"display_active"`
	PersistenceMode     string          `xml:"persistence_mode"`
	UUID                string          `xml:"uuid"`
	MinorNumber         string          `xml:"minor_number"`
	VbiosVersion        string          `xml:"vbios_version"`
	FanSpeed            string          `xml:"fan_speed"`
	PerformanceState    string          `xml:"performance_state"`
	FbMemoryUsage       fbMemoryUsage   `xml:"fb_memory_usage"`
	Bar1MemoryUsage     bar1MemoryUsage `xml:"bar1_memory_usage"`
	Temperature         temperature     `xml:"temperature"`
	PowerReadings       powerReadings   `xml:"power_readings"`
	Clocks              clocks          `xml:"clocks"`
	Processes           string          `xml:"processes"`
}

type baseInfo struct {
	Timestamp     string  `xml:"timestamp"`      // 获取信息时候的命令
	DriverVersion string  `xml:"driver_version"` // 版本号
	CudaVersion   float32 `xml:"cuda_version"`
	AttachedGpus  int     `xml:"attached_gpus"`
	Gpu           gpuInfo `xml:"gpu"`
}

type GpuInfo struct {
	NvidiaMmiLog baseInfo `xml:"nvidia_smi_log"`
}

func main() {
	//  查看命令是否存在
	// _, err := exec.LookPath("nvidia-smi")
	_, err := exec.LookPath("nvidia-smi")
	if err != nil {
		panic(err.Error())
	}
	lsCmd := exec.Command("nvidia-smi", "-x", "-q")
	output, err := lsCmd.Output()
	if err != nil {
		panic(err.Error())
	}

	var gpuInfo baseInfo
	err = xml.Unmarshal(output, &gpuInfo)
	if err != nil {
		panic(err.Error())
	}
	fmt.Printf("%+v", gpuInfo)
}
